% 1. Load data
data = readtable('附件1 数据-C题.xlsx', 'VariableNamingRule', 'preserve'); % Training set
predict_data = readtable('预测集.xlsx', 'VariableNamingRule', 'preserve'); % Prediction set
cpi_data = readtable('CPI.xlsx', 'VariableNamingRule', 'preserve'); % Yichang CPI data

% Debug: Display column names and types
disp('Training Data Columns:');
disp(data.Properties.VariableNames);
disp('Prediction Data Columns:');
disp(predict_data.Properties.VariableNames);

% Rename CPI and Year columns to avoid Chinese characters
cpi_data.Properties.VariableNames{'居民消费价格指数'} = 'CPI';
cpi_data.Properties.VariableNames{'年份'} = 'Year';

% 2. Data preprocessing
% Assume no date field; assign default year 2023 for CPI merging
data.Year = 2023 * ones(height(data), 1);
predict_data.Year = 2023 * ones(height(predict_data), 1);

% Merge Yichang CPI data
data = join(data, cpi_data, 'Keys', 'Year', 'RightVariables', 'CPI');
predict_data = join(predict_data, cpi_data, 'Keys', 'Year', 'RightVariables', 'CPI');

% Handle missing CPI values
data.CPI(isnan(data.CPI)) = mean(data.CPI, 'omitnan');
predict_data.CPI(isnan(predict_data.CPI)) = mean(predict_data.CPI, 'omitnan');

% Add other external variables (assumed values)
data.yichang_gdp = 146771 * ones(height(data), 1); % Yichang per capita GDP
data.national_unemployment = 5.2 * ones(height(data), 1); % National unemployment rate
data.job_vacancy_rate = 3.5 * ones(height(data), 1); % Job vacancy rate
data.policy_support = 80 * ones(height(data), 1); % Policy support strength
predict_data.yichang_gdp = 146771 * ones(height(predict_data), 1);
predict_data.national_unemployment = 5.2 * ones(height(predict_data), 1);
predict_data.job_vacancy_rate = 3.5 * ones(height(predict_data), 1);
predict_data.policy_support = 80 * ones(height(predict_data), 1);

% Convert categorical variables to numeric codes
data.edu_level = double(categorical(data.edu_level));
predict_data.edu_level = double(categorical(predict_data.edu_level));

% Select features and ensure all columns are numeric
features = {'age', 'edu_level', 'CPI', ...
            'yichang_gdp', 'national_unemployment', 'job_vacancy_rate', 'policy_support'};
X_train = table2array(data(:, features)); % Convert to numeric array

% Check if 'employment_status' exists, otherwise use the correct column name
if ismember('employment_status', data.Properties.VariableNames)
    y_train = data.employment_status;
elseif ismember('预测', data.Properties.VariableNames)
    y_train = data.预测;
else
    error('目标变量 "employment_status" 或 "预测" 不存在于训练数据集中。');
end

X_test = table2array(predict_data(:, features)); % Convert to numeric array

% Standardize numerical features
mu = mean(X_train, 'omitnan');
sigma = std(X_train, 0, 'omitnan');
X_train = (X_train - mu) ./ sigma;
X_test = (X_test - mu) ./ sigma;

% 3. Train XGBoost model
t = templateTree('MaxNumSplits', 20); % Limit tree depth
xgb_model = fitcensemble(X_train, y_train, ...
    'Method', 'Bag', ...
    'NumLearningCycles', 100, ...
    'Learners', t, ...
    'LearnRate', 0.1, ...
    'OptimizeHyperparameters', 'auto');

% 4. Predict on the test set
y_pred = predict(xgb_model, X_test);

% Map predictions to English labels
employment_labels = {'Not Employed', 'Employed'}; % 假设 0=Not Employed, 1=Employed
predicted_labels = employment_labels(y_pred + 1);

% Output results
disp('Prediction Results:');
for i = 1:length(predicted_labels)
    fprintf('Person %d: %s\n', i, char(predicted_labels(i)));
end

% Save prediction results
predict_data.predicted_employment_status = predicted_labels;
writetable(predict_data, 'predictions.xlsx');